#!/usr/bin/env python3

import re
import os
import sys


def read_mlir_dump_from_file(file_path):
    with open(file_path, 'r') as file:
        return file.read()


def extract_funcs(section):
    # Regular expression pattern to match function names
    pattern = r'func .*?@(\w+)\('

    # Find all matches in the input text
    function_names = re.findall(pattern, section)
    return function_names


def separate_mlir_dump(mlir_dump):
    # Split the MLIR dump into sections based on the specified patterns
    pattern = r'// -----//\s*IR Dump (Before|After)\s*([^()]*?)\s*\(.*?\)\s*//----- //(.*?)(?=// -----//\s*IR Dump (Before|After)|$)'
    matches = re.finditer(pattern, mlir_dump, re.DOTALL)

    sections = dict()
    for match in matches:
        section = match.group(0).strip()
        for func_name in extract_funcs(section):
            if func_name.startswith("__nv_"):
                continue
            if func_name in sections:
                sections[func_name].append(section)
            else:
                sections[func_name] = [section]

    return sections


def filter_sections_by_pass_name(sections, pass_name):
    filtered_sections = []
#    pass_pattern = re.compile(rf'// -----//\s*IR Dump (Before|After)\s*{re.escape(pass_name)}')

    for section in sections:
        #        if pass_pattern.search(section):
        #            filtered_sections.append(section)
        if f"// -----// IR Dump Before {pass_name}" in section or f"// -----// IR Dump After {pass_name}" in section:
            filtered_sections.append(section)

    return filtered_sections


def save(context, file_name, msg=None):
    with open(file_name, 'w') as file:
        file.write(context)

    if msg:
        print(f"Saved {msg} to {file_name}")


def save_sections_to_files(sections, output_dir):
    # Create output directory if it doesn't exist
    os.makedirs(
        output_dir, exist_ok=True)

#    for i, section in enumerate(sections):
#        file_name = f"{output_dir}/section_{i + 1}.mlir"
#        with open(file_name, 'w') as file:
#            file.write(
#                section)
#        print(f"Saved Section {i + 1} to {file_name}")
    beforeIdx = 0
    afterIdx = 0
    for section in sections:
        # Determine if the section is Before or After
        if 'IR Dump Before' in section:
            file_name = f"{output_dir}/Before_{section.split()[5]}_{beforeIdx}.mlir"
            beforeIdx = beforeIdx + 1
        else:
            file_name = f"{output_dir}/After_{section.split()[5]}_{afterIdx}.mlir"
            afterIdx = afterIdx + 1

        save(section, file_name)


def get_passes(separated_sections):
    passes = ""
    passes_changed = ""

    # Regex pattern to match the pass name
    pattern = r'// -----// IR Dump (After|Before) (\w+[_\w]*)'

    for section in separated_sections:
        first_line = section.splitlines()[0]
        pass_name_matched = re.search(pattern, first_line)
        assert pass_name_matched, f"Unknown pass name ???: {first_line} {section}"
        pass_name = pass_name_matched.group(2)  # Extract the pass name
        if f"// -----// IR Dump Before " in first_line:
            passes += pass_name + "\n"
        elif f"// -----// IR Dump After " in first_line:
            passes_changed += pass_name + "\n"
        else:
            assert False, f"Unhandled line: {first_line}"

    return passes, passes_changed


def main(dir_path):
    dir_path = os.path.dirname(dir_path)
    mlir_file_path = os.path.join(dir_path, "dump.mlir")

    # Read MLIR dump from the specified file
    mlir_dump = read_mlir_dump_from_file(
        mlir_file_path)

    # Separate the MLIR dump into sections
    separated_sections = separate_mlir_dump(
        mlir_dump)

    for func_name in separated_sections:
        output_dir = os.path.join(dir_path, func_name)
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # Extract passes
        passes, passes_changed = get_passes(separated_sections[func_name])
        # Save passes information
        save("".join(passes), os.path.join(output_dir,
             "passes.txt"), f"{func_name}_passes")
        save("".join(passes_changed), os.path.join(output_dir,
             "passes_changed.txt"), f"{func_name}_passes_changed")

        for pass_name in passes_changed.splitlines():
            # Filter sections by pass name
            filtered_sections = filter_sections_by_pass_name(
                separated_sections[func_name], pass_name)

            # Save filtered sections to files
            save_sections_to_files(
                filtered_sections, output_dir)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: python {sys.argv[0]} '<dir_path>'")
        sys.exit(1)

    dir_path = sys.argv[1]
    main(dir_path)
